package top.lywivan.tree;

/**
 * @desc
 * @author:ivan.liu
 */
public class AvlTree {

    AvlNode root;

    /**
     * 获取某一节点的树的高度
     * @param node
     * @return
     */
    public int getHeight(AvlNode node){
        return node==null?0:node.height;
    }

    /**
     * 处理RR情况，对失衡节点左旋
     * @param unbalance 失衡节点
     * @return 调整后的该子树的根节点
     */
    public AvlNode RRRotate(AvlNode unbalance){
        //找到旋转后的根节点
        AvlNode newRoot=unbalance.right;
        //处理新根节点的原来的左子树
        unbalance.right=newRoot.left;
        //处理新根节点最终的左子树
        newRoot.left=unbalance;
        //调整unbalance和newRoot的高度
        unbalance.height=Math.max(getHeight(unbalance.left),getHeight(unbalance.right))+1;
        newRoot.height=Math.max(getHeight(newRoot.left),getHeight(newRoot.right))+1;
        return newRoot;
    }

    /**
     * 处理LL情况，对失衡节点右旋
     * @param unbalance 失衡节点
     * @return 调整后的该子树的根节点
     */
    public AvlNode LLRotate(AvlNode unbalance){
        //找到旋转后的根节点
        AvlNode newRoot=unbalance.left;
        //处理新根节点的原来的左子树
        unbalance.left=newRoot.right;
        //处理新根节点最终的左子树
        newRoot.right=unbalance;
        //调整unbalance和newRoot的高度
        unbalance.height=Math.max(getHeight(unbalance.left),getHeight(unbalance.right))+1;
        newRoot.height=Math.max(getHeight(newRoot.left),getHeight(newRoot.right))+1;
        return newRoot;
    }

    /**
     * 处理LR情况，对失衡节点左子节点左旋，再对失衡节点右旋
     * @param unbalance
     * @return
     */
    public AvlNode LRRotate(AvlNode unbalance){
        //对失衡节点左子节点左旋，并将失衡节点指向新的子节点
        unbalance.left = RRRotate(unbalance.left);
        //对失衡节点右旋
        return LLRotate(unbalance);
    }

    /**
     * 处理RL情况，对失衡节点右子节点右旋，再对失衡节点左旋
     * @param unbalance
     * @return
     */
    public AvlNode RLRotate(AvlNode unbalance){
        //对失衡节点左子节点左旋，并将失衡节点指向新的子节点
        unbalance.right = LLRotate(unbalance.right);
        //对失衡节点右旋
        return RRRotate(unbalance);
    }

    /**
     * 插入操作，
     * @param key
     * @return
     */
    public boolean insert(int key){
        //递归插入key
        root=insert(root,key);
        return true;
    }

    /**
     * 递归插入key
     * @param node 当前节点，为空则进行插入
     * @param key 插入节点的值
     * @return 返回当前节点的父节点，便于平衡调整
     */
    public AvlNode insert(AvlNode node,int key){
        //terminal
        if(node==null){
            return new AvlNode(key);
        }
        //current logic
        //判断key放到左子树还是右子树
        if(key>node.key){
            //可以会旋转，导致子节点变化，因此重新设置子节点
            node.right=insert(node.right,key);
            //在更新当前节点高度前处理失衡，此时子节点的高度都已经处理完毕了
            //计算平衡因子
            if(Math.abs(getHeight(node.left)-getHeight(node.right))>1){
                //当前node是失衡节点，需要进行调整
                //判断失衡情况：要么是RR，要么是RL。此时需要判断key和node.right的大小关系即可
                if(key>node.right.key){
                    //RR的情况，更新node为最新的根节点
                    node=RRRotate(node);
                }else {
                    //RL的情况，更新node为最新的根节点
                    node=RLRotate(node);
                }
            }
        }else if(key<node.key){
            //可以会旋转，导致子节点变化，因此重新设置子节点
            node.left=insert(node.left,key);
            //在更新当前节点高度前处理失衡，此时子节点的高度都已经处理完毕了
            //计算平衡因子
            if(Math.abs(getHeight(node.left)-getHeight(node.right))>1){
                //当前node是失衡节点，需要进行调整
                //判断失衡情况：要么是LL，要么是LR。此时需要判断key和node.left的大小关系即可
                if(key<node.left.key){
                    //LL的情况，更新node为最新的根节点
                    node=LLRotate(node);
                }else {
                    //LR的情况，更新node为最新的根节点
                    node=LRRotate(node);
                }
            }
        }else {
            //key和当前节点相等，不进行操作
        }
        //递归回归时，更新当前节点高度
        node.height=Math.max(getHeight(node.left),getHeight(node.right))+1;
        //返回当前节点
        return node;
    }

    @Override
    public String toString() {
        StringBuilder sb = new StringBuilder();
        sb.append("该树的前序遍历结果为:");
        preOrder(root,sb);
        sb.append("该树的中序遍历结果为:");
        inOrder(root,sb);
        sb.append("该树的后序遍历结构为:");
        postOrder(root,sb);
        return sb.toString();
    }
    private void inOrder(AvlNode node, StringBuilder sb){
        if (node == null) {
            return;
        }
        inOrder(node.left,sb);
        sb.append(node.key).append("->");
        inOrder(node.right,sb);
    }
    private void preOrder(AvlNode node, StringBuilder sb) {
        if (node == null) {
            return;
        }
        sb.append(node.key).append("->");
        preOrder(node.left,sb);
        preOrder(node.right,sb);
    }
    private void postOrder(AvlNode node, StringBuilder sb) {
        if (node == null) {
            return;
        }
        postOrder(node.left,sb);
        postOrder(node.right,sb);
        sb.append(node.key).append("->");
    }

    public static class AvlNode {
        int key;
        AvlNode left;
        AvlNode right;
        int height=1;

        public AvlNode(int key){
            this.key=key;
        }
    }
}
